import os
import requests
import re
import json
import base64
import time
from config.development import IMAGE_GENERATION_PATH

# ========== 配置 ==========
API_KEY_DEEPSEEK = "sk-4353e122a1ca42af98ef686071a5ea20"
API_ENDPOINT_DEEPSEEK = "http://118.31.114.180:8080/openai/chat/completions"
MODEL_NAME_DEEPSEEK = "deepseek-v3"

API_URL_STABLE_DIFFUSION_TXT2IMG = "http://lyrics-vision.1730142819213311.cn-hangzhou.pai-eas.aliyuncs.com/sdapi/v1/txt2img"
API_KEY_STABLE_DIFFUSION_IMG2IMG = "http://lyrics-vision.1730142819213311.cn-hangzhou.pai-eas.aliyuncs.com/sdapi/v1/img2img"
API_KEY_STABLE_DIFFUSION = "ZTY4NTRkYmM1NmVjNmRkMzI4NGE1ODY3OGQ4YzA1MmYyZWFkZWFjMg=="

# ========== 构造提示 ==========
# 构造DeepSeek的提示，要求用户分析歌词，设计视觉叙事
def build_prompt(lyrics: str, style: str) -> str:
    return f"""
你是一个动漫导演兼分镜师，请按如下要求进行创作：

【第一步】
你需要仔细理解歌词的意义，分析歌词中所讲述的情感线索、时间顺序、人物关系和核心主题。

【第二步】
基于你对歌词的理解，以{style}风格创建，设计一段完整的视觉叙事，用 4 个动漫风格的分镜画面表现出来。这些分镜要构成一个有逻辑顺序的故事，不是孤立的画面。注意提示词要能让生图模型生成动漫风格的图片。

歌词内容:
        {lyrics}

【第三步】
请按照以下JSON格式输出:

```json
        [
          {{
            "Title": "分镜标题（简短有力，概括画面主题）",
            "Plot": "此画面所承载的情节信息",
            "English": "用于图像生成的详细英文描述（anime style，给出关键的场景人物和情节，要能让模型很好地理解）",
            "Chinese": "中文文艺描述，用于展示阅读"
          }},
          ...
        ]
        只返回JSON格式内容，不要有其他解释。
```
【注意】
- 英文描述用于图像生成，务必详细，视觉清晰，有人物、动作、环境、情绪
- 中文描述用于辅助理解，不用于生成
- 整体风格应为 anime style，并保持4个画面角色、色调、构图一致
"""

# ========== 请求 DeepSeek ==========
# 向DeepSeek API发送请求，返回生成的回复
def query_deepseek(prompt: str) -> str:
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY_DEEPSEEK}"
    }
    data = {
        "model": MODEL_NAME_DEEPSEEK,
        "messages": [
            {"role": "user", "content": prompt}
        ]
    }
    response = requests.post(API_ENDPOINT_DEEPSEEK, json=data, headers=headers)
    response.raise_for_status()
    return response.json()['choices'][0]['message']['content']

# ========== 提取分镜 ==========
# 解析DeepSeek的回复，提取出4个分镜的标题、剧情概览、英文描述、中文描述
def extract_bilingual_scenes(response_text: str):
    scenes = json.loads(response_text)
    extracted_scenes = []
    for scene in scenes:
        extracted_scenes.append({
            "Title": scene["Title"],
            "Plot": scene["Plot"],
            "English": scene["English"],
            "Chinese": scene["Chinese"]
        })
    return extracted_scenes

# ========== 生成图片 ==========
#传入分镜标号参数，生成对应分镜的图片
def generate_image(positive_prompt=None, negative_prompt=None, num_images=1,scene_index=0):
    payload = {
        "enable_hr": False,
        "denoising_strength": 0,
        "firstphase_width": 0,
        "firstphase_height": 0,
        "batch_size": num_images,
        "prompt": positive_prompt if positive_prompt else "a beautiful landscape with mountains and lakes, sunset, highly detailed, 4k",
        "negative_prompt": negative_prompt if negative_prompt else "blurry, low quality, distorted",
        "styles": [],
        "seed": -1,
        "steps": 20,
        "width": 512,
        "height": 512,
        "cfg_scale": 7,
        "sampler_name": "Euler a",
        "n_iter": 1,
        "restore_faces": False,
        "tiling": False
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY_STABLE_DIFFUSION}"
    }
    try:
        response = requests.post(API_URL_STABLE_DIFFUSION_TXT2IMG, headers=headers, data=json.dumps(payload))
        if response.status_code == 200:
            result = response.json()
            for i, image_data in enumerate(result.get("images", [])):
                image_bytes = base64.b64decode(image_data)
                # 生成文件名
                file_name = f"generated_image_{scene_index}_{i}_{int(time.time())}.png"

                with open(f"{IMAGE_GENERATION_PATH}\\{file_name}", "wb") as f:
                    f.write(image_bytes)
                print(f"图像已保存为 {IMAGE_GENERATION_PATH}\{file_name}")
        else:
            print(f"请求失败，状态码: {response.status_code}")
            print(response.text)
    except Exception as e:
        print(f"发生错误: {str(e)}")

    return f"{IMAGE_GENERATION_PATH}\{file_name}"

# ========== 修改图片 ==============
# 传入分镜标号参数，修改对应分镜的图片
def modify_image(positive_prompt=None, negative_prompt=None,scene_index=0):
    current_working_directory = os.getcwd()
    # 初始化一个空列表用于存储base64编码的图像数据
    image_base64_list = []
    with open(f"{current_working_directory}/generated_image_{scene_index}_0.png", "rb") as f:
        image_bytes = f.read()
        #转化为base64编码
        image_base64 = base64.b64encode(image_bytes).decode()
        # 将base64字符串添加到列表中
        image_base64_list.append(image_base64)
    #调用接口
    payload = {
        "prompt": positive_prompt if positive_prompt else "a beautiful landscape with mountains and lakes, sunset, highly detailed, 4k",
        "negative_prompt": negative_prompt if negative_prompt else "blurry, low quality, distorted",
        "styles": [],
        "seed": -1,
        "sampler_name": "Euler a",
        "batch_size": 1,
        "n_iter": 1,
        "steps": 20,
        "cfg_scale": 7,
        "width": 512,
        "height": 512,
        "denoising_strength": 0.75,
        "init_images": image_base64_list,
    }
    headers = {
        "Content-Type": "application/json",
        "Authorization": f"Bearer {API_KEY_STABLE_DIFFUSION}"
    }
    try:
        response = requests.post(API_KEY_STABLE_DIFFUSION_IMG2IMG, headers=headers, data=json.dumps(payload))
        if response.status_code == 200:
            result = response.json()
            for i, image_data in enumerate(result.get("images", [])):
                image_bytes = base64.b64decode(image_data)
                # 生成唯一文件名
                file_name = f"modified_image_{scene_index}_{i}_{int(time.time())}.png"
                with open(f"{current_working_directory}\{file_name}", "wb") as f:
                    f.write(image_bytes)
        else:
            print(f"请求失败，状态码: {response.status_code}")
            print(response.text)
    except Exception as e:
        print(f"发生错误: {str(e)}")
    
    return f"{current_working_directory}\{file_name}"
# ========== 主程序 ==========
def main():
    #打开歌词文件
    with open("lyrics.txt", "r", encoding="utf-8") as f:
         lyrics = f.read()
    #lyrics = input("请输入歌词（将为其生成4个动漫风格的分镜）：\n")
    prompt = build_prompt(lyrics,"美丽")
    print("\n🧠 正在生成分镜，请稍候...\n")
    try:
        response_text = query_deepseek(prompt)
    except Exception as e:
        print("❌ 请求出错：", e)
        return
    print("🎬 模型原始回复：\n")
    print(response_text)
    scenes = extract_bilingual_scenes(response_text)

    print("\n🎨 提取后的分镜结果：\n")
    scene_index = 0
    for scene in scenes:
        print(f"Scene 标题: {scene['title']}")
        print(f"剧情概览: {scene['plot']}")
        print(f"Prompt for image gen (EN):\n{scene['english']}\n")
        print(f"Prompt for image gen (ZH):\n{scene['chinese']}\n")
        generate_image(positive_prompt=scene['english'],scene_index=scene_index)
        #这里插入修改描述的语句
        modify_image(positive_prompt=scene['english']+"more beautiful and more detaild",scene_index=scene_index)
        scene_index += 1
if __name__ == "__main__":
    main()